from argparse import ArgumentParser
from project_utils.image_io import *

import matplotlib.pyplot as plt
import numpy as np
import os.path as osp


def paddedsize(AB, CD=None, PARAM=None):
    """Computes padded sizes useful for FFT-based filtering.

    `AB`: A two-element size vector.
    `CD`: A two-element size vector.
    `PARAM`: Extra widening parameter.

    ```
    # PQ: 2 * AB
    PQ = paddedsize(AB)

    # PQ[0], PQ[1]: 2 ^ nextpow2(2 * m), where m is max(AB)
    PQ = paddedsize(AB, PARAM='PWR2')

    # PQ: the smallest even integers greater than or equal to AB + CD - 1
    PQ = paddedsize(AB, CD)

    # PQ[0], PQ[1]: 2 ^ nextpow2(2 * m), where m is max([AB, CD])
    PQ = paddedsize(AB, CD, 'PWR2')
    ```
    """

    AB = np.array(AB)

    if CD is not None:
        CD = np.array(CD)

    if PARAM is None:
        m = AB if CD is None else np.ceil((AB + CD - 1) / 2)
        PQ = 2 * m
    elif PARAM == "PWR2":
        m = max(AB) if CD is None else max([AB, CD])
        P = 2 ** np.ceil(np.log2(2 * m))
        PQ = [P, P]
    else:
        raise ValueError("Wrong value for PARAM")

    return PQ


def freqz2(a, shape):
    """Obtain a two-dimensional filter in frequency domain from a spatial domain filter.

    `a`: Spatial domain filter, such as 'sodel'.
    `shape`: Size of the filter to be generated, in shape of (Rows, Columns).
    """

    freq = np.fft.fft2(a, s=shape)
    return np.abs(np.fft.fftshift(freq))


def lpfilter(type, M, N, D0, n=1.0):
    """Obtain different types of low-pass filters.

    `type`: filter type, choose from ('ideal', 'btw', 'gaussian').
    `M`: the number of rows of the filter.
    `N`: the number of columns of the filter.
    `D0`: cut-off frequency (a positive number).
    `n`: the order of the 'btw' filter.

    Returns the transfer function `H` and distance vector `D` of the low-pass filter.
    """

    u = np.arange(0, M)
    v = np.arange(0, N)

    idx = [i for i in range(len(u)) if u[i] > M // 2]
    u[idx] = [u[i] - M for i in idx]
    idy = [i for i in range(len(v)) if v[i] > N // 2]
    v[idy] = [v[i] - N for i in idy]

    U, V = np.meshgrid(v, u)

    # Calculate the distance vector D, centered around the boundary
    D = np.sqrt(U**2 + V**2)

    # Perform different calculations based on filter type
    if type == "ideal":
        H = np.array(D <= D0, dtype=float)
    elif type == "btw":
        H = 1 / (1 + (D / D0) ** (2 * n))
    elif type == "gaussian":
        H = np.exp(-(D**2) / (2 * (D0**2)))
    else:
        raise ValueError("Unknown filter type.")

    return H, D


def hpfilter(type, M, N, D0, n=1.0):
    """Obtain different types of high-pass filters.

    `type`: filter type, choose from ('ideal', 'btw', 'gaussian').
    `M`: the number of rows of the filter.
    `N`: the number of columns of the filter.
    `D0`: cut-off frequency (a positive number).
    `n`: the order of the 'btw' filter.

    Returns the transfer function `H` of the high-pass filter.
    """

    H, D = lpfilter(type, M, N, D0, n)

    return 1 - H


def dftfilt2(f, H):
    """Frequency domain filtering processing.

    `f`: the image to be filtered.
    `H`: the filter in frequency domain.

    Returns the filtered image.
    """

    s = (H.shape[0], H.shape[1])

    # Obtain the FFT of the padded input
    F = np.fft.fft2(f.astype(np.float32), s)

    # Perform filtering
    g = np.real(np.fft.ifft2(H.astype(np.complex64) * F))

    # Crop to original size
    g = g[0 : f.shape[0], 0 : f.shape[1]]

    return g.astype(np.uint8)


def show_filter_3d(H, B):
    """Obtain the 3d image of the the filter.

    `H`: two-dimensional filter in frequency domain.
    `B`: the title of the plot.
    """

    # Create grid coordinates
    M, N = H.shape
    u = np.arange(-M // 2, M // 2)
    v = np.arange(-N // 2, N // 2)
    u, v = np.meshgrid(u, v, indexing="ij")

    # Shift the filter to center
    H_shifted = np.fft.fftshift(H)

    # Create 3D plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    # Plot the surface
    ax.plot_surface(u, v, H_shifted, cmap="viridis")

    # Set labels and title
    ax.set_xlabel("u")
    ax.set_ylabel("v")
    ax.set_zlabel("Magnitude")
    ax.set_title(B)

    # Show the plot
    plt.show()


def Frequency_Filtering_test(
    image_path, Filter_type, Filter_model, Cutoff_coefficient, Filter_name, save_name
):
    """Frequency filtering test function.

    `image_path`: filepath of the source image.
    `Filter_type`: select filter from ('ideal', 'btw', 'gaussian').
    `Filter_model`: choose between low-pass filter and high-pass filter.
    `Cutoff_coefficient`: cutoff frequency used in the filtering.
    `Filter_name`: title of the 3D plot of the filter.
    `save_name`: filepath of the target image.
    """

    assert Filter_model in [
        "LOW",
        "HIGH",
    ], "Choose between low pass filters and high pass filters."

    gray_image = load_image_as_grayscale(image_path, "rgb")
    image_size = (gray_image.shape[0], gray_image.shape[1])
    fill_size = paddedsize(image_size)

    if Filter_model == "LOW":
        filter, vector_matrix = lpfilter(
            Filter_type,
            fill_size[0],
            fill_size[1],
            D0=Cutoff_coefficient * fill_size[0],
            n=1,
        )

    if Filter_model == "HIGH":
        filter = hpfilter(
            Filter_type,
            fill_size[0],
            fill_size[1],
            D0=Cutoff_coefficient * fill_size[0],
            n=1,
        )

    show_filter_3d(filter, Filter_name)

    filter_image = dftfilt2(gray_image, filter)
    save_grayscale_image(save_name, filter_image, False)

    plt.imshow(filter_image, cmap="gray")
    plt.show()


def main_procedure(cmdargs):
    source_path = osp.abspath(cmdargs.source)
    target_path = osp.abspath(cmdargs.target)

    print(cmdargs)

    default_title = f"{cmdargs.filter} {cmdargs.bandpass} pass filter".lower()
    plot_title = cmdargs.plot_title or default_title
    default_freq = 0.05 if cmdargs.bandpass == "LOW" else 0.5
    cutoff_freq = cmdargs.cutoff_freq or default_freq

    Frequency_Filtering_test(
        image_path=source_path,
        Filter_type=cmdargs.filter,
        Filter_model=cmdargs.bandpass,
        Cutoff_coefficient=cutoff_freq,
        Filter_name=plot_title,
        save_name=target_path,
    )


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Demonstration of the Effects of Different Filters."
    )

    parser.add_argument(
        "filter",
        metavar="FILTER",
        choices=["ideal", "gaussian", "btw"],
        help="Select the type of the filter in the demonstration.",
    )
    parser.add_argument(
        "bandpass",
        metavar="BANDPASS",
        choices=["LOW", "HIGH"],
        help="Select band pass/reject strategy.",
    )

    parser.add_argument(
        "--source", type=str, required=True, help="Filepath of the source image."
    )
    parser.add_argument(
        "--target", type=str, required=True, help="Filepath of the target image."
    )
    parser.add_argument(
        "--plot_title", type=str, help="Title of the 3D plot of the filter."
    )
    parser.add_argument(
        "--cutoff_freq", type=float, help="Cutoff frequency used in the filtering."
    )

    main_procedure(parser.parse_args())
